package org.xp.judgeplatform.util;

import io.jsonwebtoken.*;
import io.jsonwebtoken.io.Decoders;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.stereotype.Component;

import javax.crypto.SecretKey;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

@Component
@Slf4j
public class JwtUtil {

    @Value("${jwt.secret}")
    private String secret;

    @Value("${jwt.expiration}")
    private Long expiration;

    @Value("${jwt.refresh-expiration}")
    private Long refreshExpiration;

    private SecretKey secretKey;
    private JwtParser jwtParser;

    @PostConstruct
    public void init() {
        this.secretKey = Keys.hmacShaKeyFor(Decoders.BASE64.decode(secret));
        this.jwtParser = Jwts.parser().verifyWith(this.secretKey).build();
        log.info("JWT工具初始化完成");
    }

    private static final String CLAIM_KEY_USERNAME = "username";
    private static final String CLAIM_KEY_TYPE = "type";
    private static final String TOKEN_TYPE_REFRESH = "refresh";

    public String generateToken(UserDetails userDetails) {
        Map<String, Object> claims = new HashMap<>();
        claims.put(CLAIM_KEY_USERNAME, userDetails.getUsername());
        return createToken(claims, userDetails.getUsername(), expiration);
    }

    public String generateRefreshToken(UserDetails userDetails) {
        Map<String, Object> claims = new HashMap<>();
        claims.put(CLAIM_KEY_USERNAME, userDetails.getUsername());
        claims.put(CLAIM_KEY_TYPE, TOKEN_TYPE_REFRESH);
        return createToken(claims, userDetails.getUsername(), refreshExpiration);
    }

    private String createToken(Map<String, Object> claims, String subject, Long expiration) {
        Date now = new Date();
        Date expiryDate = new Date(now.getTime() + expiration * 1000);

        return Jwts.builder()
                .claims(claims)
                .subject(subject)
                .issuedAt(now)
                .expiration(expiryDate)
                .signWith(secretKey)
                .compact();
    }

    public String getUsernameFromToken(String token) {
        return getClaimFromToken(token, Claims::getSubject);
    }

    public Date getExpirationDateFromToken(String token) {
        return getClaimFromToken(token, Claims::getExpiration);
    }

    public <T> T getClaimFromToken(String token, Function<Claims, T> claimsResolver) {
        final Claims claims = getAllClaimsFromToken(token);
        return claimsResolver.apply(claims);
    }

    private Claims getAllClaimsFromToken(String token) {
        try {
            return jwtParser.parseSignedClaims(token).getPayload();
        } catch (ExpiredJwtException e) {
            log.warn("JWT令牌已过期: {}", e.getMessage());
            throw e;
        } catch (SignatureException e) {
            log.error("JWT签名无效: {}", e.getMessage());
            throw e;
        } catch (MalformedJwtException e) {
            log.error("JWT格式错误: {}", e.getMessage());
            throw e;
        } catch (JwtException e) {
            log.error("解析JWT令牌失败: {}", e.getMessage());
            throw e;
        }
    }

    public Boolean isTokenExpired(String token) {
        try {
            final Date expiration = getExpirationDateFromToken(token);
            return expiration.before(new Date());
        } catch (JwtException e) {
            log.warn("检查令牌过期时发生异常: {}", e.getMessage());
            return true;
        }
    }

    public Boolean validateAccessToken(String token, UserDetails userDetails) {
        try {
            final String username = getUsernameFromToken(token);
            final String tokenType = getClaimFromToken(token,
                    claims -> claims.get(CLAIM_KEY_TYPE, String.class));

            if (TOKEN_TYPE_REFRESH.equals(tokenType)) {
                return false;
            }

            return username.equals(userDetails.getUsername()) && !isTokenExpired(token);
        } catch (JwtException e) {
            log.warn("访问令牌验证失败: {}", e.getMessage());
            return false;
        }
    }

    public Boolean validateRefreshToken(String token, UserDetails userDetails) {
        try {
            final String username = getUsernameFromToken(token);
            final String tokenType = getClaimFromToken(token,
                    claims -> claims.get(CLAIM_KEY_TYPE, String.class));

            if (!TOKEN_TYPE_REFRESH.equals(tokenType)) {
                return false;
            }

            return username.equals(userDetails.getUsername()) && !isTokenExpired(token);
        } catch (JwtException e) {
            log.warn("刷新令牌验证失败: {}", e.getMessage());
            return false;
        }
    }

    public Boolean canTokenBeRefreshed(String token) {
        return !isTokenExpired(token);
    }

    public String refreshToken(String token) {
        try {
            final Claims claims = getAllClaimsFromToken(token);

            String tokenType = claims.get(CLAIM_KEY_TYPE, String.class);
            if (!TOKEN_TYPE_REFRESH.equals(tokenType)) {
                throw new JwtException("只能使用刷新令牌进行刷新");
            }

            return Jwts.builder()
                    .claims(claims)
                    .issuedAt(new Date())
                    .expiration(new Date(System.currentTimeMillis() + expiration * 1000))
                    .signWith(secretKey)
                    .compact();
        } catch (JwtException e) {
            log.error("刷新JWT令牌失败: {}", e.getMessage());
            throw e;
        }
    }

    public String extractTokenFromHeader(String authHeader) {
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            return authHeader.substring(7);
        }
        return null;
    }
}